"""
MCP消息处理器

负责消息的创建、序列化、反序列化和验证
"""

from typing import Dict, Any, Optional
from datetime import datetime
import uuid


class MCPMessageHandler:
    """
    MCP消息处理器
    
    功能：
    1. 创建各类MCP消息（CONNECT、REQUEST、PING等）
    2. 验证消息格式
    3. 消息序列化/反序列化
    
    消息格式规范：
    {
        "type": "MESSAGE_TYPE",        # 消息类型
        "message_id": "uuid",          # 消息唯一标识
        "session_id": "uuid",          # 会话ID
        "payload": {...},              # 负载数据
        "timestamp": "2025-10-22T10:00:00Z"  # 时间戳
    }
    """
    
    @staticmethod
    def create_connect_message(session_id: str) -> Dict[str, Any]:
        """
        创建CONNECT消息
        
        用于建立MCP连接时发送的第一条消息
        
        Args:
            session_id: 会话ID
            
        Returns:
            Dict[str, Any]: CONNECT消息
        """
        return {
            "type": "CONNECT",
            "message_id": str(uuid.uuid4()),
            "session_id": session_id,
            "payload": {
                "protocol_version": "1.0",
                "client_info": {
                    "name": "fengxuanLLModel-mcp-client",
                    "version": "1.0.0",
                    "platform": "python"
                }
            },
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }
    
    @staticmethod
    def create_request_message(
        session_id: str,
        message_id: str,
        payload: Dict[str, Any]
    ) -> Dict[str, Any]:
        """
        创建REQUEST消息
        
        用于向MCP服务器发送模型请求
        
        Args:
            session_id: 会话ID
            message_id: 消息ID
            payload: 请求负载数据
            
        Returns:
            Dict[str, Any]: REQUEST消息
        """
        return {
            "type": "REQUEST",
            "message_id": message_id,
            "session_id": session_id,
            "payload": payload,
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }
    
    @staticmethod
    def create_ping_message(session_id: str) -> Dict[str, Any]:
        """
        创建PING消息
        
        用于心跳检测
        
        Args:
            session_id: 会话ID
            
        Returns:
            Dict[str, Any]: PING消息
        """
        return {
            "type": "PING",
            "message_id": str(uuid.uuid4()),
            "session_id": session_id,
            "payload": {},
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }
    
    @staticmethod
    def create_pong_message(session_id: str) -> Dict[str, Any]:
        """
        创建PONG消息
        
        用于响应PING心跳
        
        Args:
            session_id: 会话ID
            
        Returns:
            Dict[str, Any]: PONG消息
        """
        return {
            "type": "PONG",
            "message_id": str(uuid.uuid4()),
            "session_id": session_id,
            "payload": {},
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }
    
    @staticmethod
    def create_disconnect_message(session_id: str) -> Dict[str, Any]:
        """
        创建DISCONNECT消息
        
        用于主动断开连接
        
        Args:
            session_id: 会话ID
            
        Returns:
            Dict[str, Any]: DISCONNECT消息
        """
        return {
            "type": "DISCONNECT",
            "message_id": str(uuid.uuid4()),
            "session_id": session_id,
            "payload": {
                "reason": "client_initiated"
            },
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }
    
    @staticmethod
    def is_connect_ack(message: Dict[str, Any]) -> bool:
        """验证是否为CONNECT_ACK消息"""
        return message.get("type") == "CONNECT_ACK"
    
    @staticmethod
    def is_response(message: Dict[str, Any]) -> bool:
        """验证是否为RESPONSE消息"""
        return message.get("type") == "RESPONSE"
    
    @staticmethod
    def is_error(message: Dict[str, Any]) -> bool:
        """验证是否为ERROR消息"""
        return message.get("type") == "ERROR"
    
    @staticmethod
    def is_ping(message: Dict[str, Any]) -> bool:
        """验证是否为PING消息"""
        return message.get("type") == "PING"
    
    @staticmethod
    def is_pong(message: Dict[str, Any]) -> bool:
        """验证是否为PONG消息"""
        return message.get("type") == "PONG"
    
    @staticmethod
    def validate_message(message: Dict[str, Any]) -> bool:
        """
        验证消息格式是否符合MCP规范
        
        Args:
            message: 待验证的消息
            
        Returns:
            bool: 是否有效
        """
        required_fields = ["type", "message_id", "session_id", "payload", "timestamp"]
        return all(field in message for field in required_fields)
    
    @staticmethod
    def extract_error(message: Dict[str, Any]) -> tuple[str, str]:
        """
        从ERROR消息中提取错误信息
        
        Args:
            message: ERROR消息
            
        Returns:
            tuple[str, str]: (错误码, 错误信息)
        """
        payload = message.get("payload", {})
        error_code = payload.get("code", "UNKNOWN")
        error_message = payload.get("message", "未知错误")
        return error_code, error_message

